#include "DGGSVisualizer.hpp"
#include <GridTIFF.hpp>
#include <DGGSFeatureLayer.hpp>
#include <DGGSFeature.hpp>
#include <DGGSFeatureClass.hpp>
#include <cmath>
#include <iostream>
#include <algorithm>
#include <numeric>
#include <chrono>
#include <iomanip>

_NNU_DGGS_BEGIN
namespace DGGSViewer
{
    // ThreadPool实现
    ThreadPool::ThreadPool(size_t threads) : stop(false)
    {
        for (size_t i = 0; i < threads; ++i)
            workers.emplace_back([this]
            {
                while (true)
                {
                    std::function<void()> task;
                    {
                        std::unique_lock<std::mutex> lock(this->queue_mutex);
                        this->condition.wait(lock, [this]
                        {
                            return this->stop || !this->tasks.empty();
                        });
                        if (this->stop && this->tasks.empty())
                            return;
                        task = std::move(this->tasks.front());
                        this->tasks.pop();
                    }
                    task();
                }
            });
    }

    template <class F, class... Args>
    auto ThreadPool::enqueue(F&& f, Args&&... args)
        -> std::future<typename std::result_of<F(Args...)>::type>
    {
        using return_type = typename std::result_of<F(Args...)>::type;

        auto task = std::make_shared<std::packaged_task<return_type()>>(
            std::bind(std::forward<F>(f), std::forward<Args>(args)...));

        std::future<return_type> res = task->get_future();
        {
            std::unique_lock<std::mutex> lock(queue_mutex);
            if (stop)
                throw std::runtime_error("enqueue on stopped ThreadPool");
            tasks.emplace([task]()
            {
                (*task)();
            });
        }
        condition.notify_one();
        return res;
    }

    ThreadPool::~ThreadPool()
    {
        {
            std::unique_lock<std::mutex> lock(queue_mutex);
            stop = true;
        }
        condition.notify_all();
        for (std::thread& worker : workers)
            worker.join();
    }

    void DGGSVisualizer::tiffHandler(const std::string& filePath)
    {
        size_t threadCount = std::max(1u, std::thread::hardware_concurrency() - 1);
        m_threadPool = std::make_unique<ThreadPool>(threadCount);
        std::cout << "创建线程池，线程数: " << threadCount << std::endl;

        const auto gridTiff = new GridTIFF();
        gridTiff->loadFromFile(filePath.c_str());

        m_kernel = CreateKernelObj(gridTiff->getGridType());

        type = gridTiff->getElementType();

        const int n = gridTiff->getSize(); // n代表行数和列数
        const int totalSize = n * n; // 总数据量是n²

        std::cout << "文件数据大小: " << totalSize << " 个单元" << std::endl;

        const GridTransform gridTrans = gridTiff->getTrans();
        const int gridSize = gridTiff->getSize();
        const DGGridBBox gridBBox = getGridBoxFromGridTrans(gridTrans, gridSize);

        // 按批次处理数据，使用线程池
        const size_t numBatches = (totalSize + m_batchSize - 1) / m_batchSize;
        std::vector<std::future<std::vector<Cell>>> futures;
        futures.reserve(numBatches);

        std::cout << "开始并行处理，批次数量: " << numBatches << std::endl;

        for (size_t i = 0; i < numBatches; ++i)
        {
            // 计算当前批次的索引范围
            const size_t startIdx = i * m_batchSize;
            const size_t endIdx = std::min(startIdx + m_batchSize, static_cast<size_t>(totalSize));

            // 创建索引向量
            std::vector<int> indices(endIdx - startIdx);
            std::iota(indices.begin(), indices.end(), static_cast<int>(startIdx));

            // 将任务添加到线程池
            futures.push_back(m_threadPool->enqueue(
                [this, gridTiff, indices, gridBBox]()
                {
                    std::vector<Cell> batchCells;
                    this->processTiffBatch(gridTiff, indices, gridBBox, batchCells);
                    return batchCells;
                }
            ));
        }

        // 收集所有批次处理的结果
        size_t totalProcessed = 0;

        for (auto& future : futures)
        {
            auto batchCells = future.get();
            totalProcessed += batchCells.size();

            // 将批次结果合并到主cells容器
            std::lock_guard lock(m_cellsMutex);
            cells.insert(cells.end(), batchCells.begin(), batchCells.end());
        }

        std::cout << "并行处理完成，总处理单元数: " << totalProcessed << std::endl;

        // 清理资源
        delete gridTiff;
        m_threadPool.reset(); // 销毁线程池

        DestroyKernelObj(m_kernel);
    }

    void DGGSVisualizer::jsonHandler(const std::string& filePath)
    {
        const auto dggsFeature = new DGGSFeatureLayer();
        dggsFeature->loadFromFile(filePath.c_str());

        const auto featureClass = dggsFeature->getFeatureClass();

        m_kernel = CreateKernelObj(featureClass->getGridType());
        type = featureClass->getElementType();

        for (int i = 0; i < dggsFeature->getFeatureCount(); i++)
        {
            const auto feature = dggsFeature->getFeature(i);
            const auto geometry = feature->getGeometry();
            for (int j = 0; j < geometry->getGridCount(); j++)
            {
                Cell cell;

                const auto code = geometry->getGridCodeByIndex(j);

                if (type == DGGSElementType::Vertex)
                {
                    GeoCoord geo;
                    m_kernel->decode(code, geo);
                    cell.m_coordinates.emplace_back(geo);
                }
                else
                {
                    DGGSElementCodeIter* vertexIter = nullptr;
                    int vertexCount = 0;

                    m_kernel->query(DGGSTopoType::RelativeVertexes, code, vertexIter, vertexCount);

                    for (int k = 0; k < vertexCount; k++)
                    {
                        GeoCoord geo;
                        m_kernel->decode(vertexIter[k].code, geo);
                        cell.m_coordinates.emplace_back(geo);
                    }

                    delete[] vertexIter;
                }

                cell.values.emplace_back(i);

                cells.emplace_back(cell);
            }
        }

        DestroyKernelObj(m_kernel);
    }

    void DGGSVisualizer::processTiffBatch(GridTIFF* gridTiff, const std::vector<int>& indices,
                                          const DGGridBBox& gridBBox, std::vector<Cell>& result) const
    {
        // 预分配结果容器空间，提高性能
        result.reserve(indices.size());

        // 获取波段数和数据类型
        const auto bandCount = gridTiff->getBandCount();
        const DGGSDataType dataType = gridTiff->getDataType();

        // 预先获取所有波段的数据 (只获取一次，避免重复调用)
        std::vector<void*> bandValues(bandCount);
        for (int bandIdx = 0; bandIdx < bandCount; bandIdx++)
        {
            bandValues[bandIdx] = gridTiff->getBandValue(bandIdx);
        }

        const int n = gridTiff->getSize() ; // 行列数

        // 处理每个索引
        for (const int idx : indices)
        {
            // 计算行列号
            const int row = idx / n + gridBBox.rowMin;
            const int col = idx % n + gridBBox.colMax - gridBBox.colSize + 1;

            Cell cell;

            DGGSElementCode code;
            code.level = gridTiff->getTrans().level;
            code.basePartition = gridTiff->getTrans().basePartition;
            code.elementType = type;
            code.row=row;
            code.col=col;


            if (type == DGGSElementType::Vertex)
            {
                GeoCoord geo;
                m_kernel->decode(code, geo);
                cell.m_coordinates.emplace_back(geo);
            }
            else
            {
                DGGSElementCodeIter* vertexIter = nullptr;
                int vertexCount = 0;

                m_kernel->query(DGGSTopoType::RelativeVertexes, code, vertexIter, vertexCount);

                for (int i = 0; i < vertexCount; i++)
                {
                    GeoCoord geo;
                    m_kernel->decode(vertexIter[i].code, geo);
                    cell.m_coordinates.emplace_back(geo);
                }

                delete[] vertexIter;
            }

            // 预分配波段值的空间
            cell.values.resize(bandCount);

            // 提前计算idx的偏移量，避免重复计算
            //const size_t byteOffset = idx;

            int x,y;
            this->Code2XY(code,x,y,gridBBox);
            size_t xSize=(size_t) gridBBox.rowSize*2;
            size_t ySize=(size_t) gridBBox.colSize*2;
            if(x==-1||y==-1)
                continue;
            const size_t byteOffset = (ySize - 2 - y) * xSize + x;



            // 遍历每个波段
            for (int bandIdx = 0; bandIdx < bandCount; bandIdx++)
            {
                const void* bandValue = bandValues[bandIdx];
                double value = 0.0;

                // 根据数据类型读取值
                switch (dataType)
                {
                case DGGS_Byte:
                    value = static_cast<double>(static_cast<const unsigned char*>(bandValue)[byteOffset]);
                    break;
                case DGGS_UInt16:
                    value = static_cast<double>(static_cast<const unsigned short*>(bandValue)[byteOffset]);
                    break;
                case DGGS_Int16:
                    value = static_cast<double>(static_cast<const short*>(bandValue)[byteOffset]);
                    break;
                case DGGS_UInt32:
                    value = static_cast<double>(static_cast<const unsigned int*>(bandValue)[byteOffset]);
                    break;
                case DGGS_Int32:
                    value = static_cast<double>(static_cast<const int*>(bandValue)[byteOffset]);
                    break;
                case DGGS_Float32:
                    value = static_cast<double>(static_cast<const float*>(bandValue)[byteOffset]);
                    break;
                case DGGS_Float64:
                    value = static_cast<double>(static_cast<const double*>(bandValue)[byteOffset]);
                    break;
                default:
                    value = 0.0;
                    break;
                }

                // 存储波段值
                cell.values[bandIdx] = value;
            }

            result.emplace_back(std::move(cell));
        }
    }

    bool DGGSVisualizer::XY2Code(int x, int y, DGGSElementCode& code, const DGGridBBox& gridBBox)
    {
        const int level = code.level;
        //保证格网的编码正确
        if (const long long int maxRow = gridBBox.rowMin + gridBBox.rowSize + 1; static_cast<long long>(y) - x >= 0 -
            maxRow && static_cast<long long>(y) - x <= maxRow - 1 && static_cast<long long>(y) + x >= maxRow - 1 &&
            static_cast<long long>(y) + x <= maxRow * 3 - 2)
        {
            if (x % 2 == 0 && y % 2 == 0)
                x--;
            if (x % 2 != 0 && y % 2 != 0)
                x--;

            constexpr int x0 = 0;
            const int y0 = static_cast<int>(maxRow) - 1;
            const int row = (y - y0 + (x - x0)) / 2;
            const int col = (x - x0 - (y - y0)) / 2;
            code.row = row;
            code.col = col;
            code.level = level;
            return true;
        }

        return false;
    }

    bool DGGSVisualizer::Code2XY(DGGSElementCode code, int& x, int& y, DGGridBBox gridBox)
    {
        if (code.basePartition != gridBox.basePartition)
        {
            //表示没找到，可以跳过
            x = -1, y = -1;
            return false;
        }

        long long int colMin = gridBox.colMax - gridBox.colSize + 1;
        long long int rowMin = gridBox.rowMin;
        int x0 = 0;
        int y0 = (int)gridBox.rowSize - 1;

        x = int(code.row - rowMin + code.col - colMin + x0);
        y = int((code.col - colMin) - (code.row - rowMin) + y0);
        if (y < 0 || x < 0)
        {
            //表示没找到，可以跳过
            x = -1, y = -1;
            return false;
        }
        return true;
    }

    DGGridBBox DGGSVisualizer::getGridBoxFromGridTrans(const GridTransform gridTrans, const int gridSize)
    {
        DGGridBBox gridBBox;
        gridBBox.basePartition = gridTrans.basePartition;
        gridBBox.colSize = gridSize;
        gridBBox.rowSize = gridSize;
        gridBBox.rowMin = gridTrans.rowOrigin;
        gridBBox.colMax = gridTrans.colOrigin + gridBBox.colSize - 1;

        return gridBBox;
    }

    void DGGSVisualizer::loadData(const std::string& filePath)
    {
        // 清空cells
        cells.clear();

        // 获取文件后缀
        std::string extension;
        if (const size_t pos = filePath.find_last_of('.'); pos != std::string::npos)
        {
            extension = filePath.substr(pos);
        }

        if (extension == ".json")
        {
            jsonHandler(filePath);
        }
        else if (extension == ".geojson")
        {
            jsonHandler(filePath);
        }
        else if (extension == ".tiff" || extension == ".tif")
        {
            // 处理TIFF文件
            tiffHandler(filePath);
        }
        else
        {
            // 不支持的文件类型
            std::cout << "不支持的文件类型: " << extension << std::endl;
        }
    }

    void DGGSVisualizer::loadMultipleData(const std::vector<std::string>& filePaths)
    {
        // 清空cells
        cells.clear();

        if (filePaths.empty())
        {
            std::cout << "没有文件需要加载" << std::endl;
            return;
        }

        // 保存第一个有效文件的元素类型
        bool typeSet = false;

        // 依次加载每个文件
        for (size_t i = 0; i < filePaths.size(); ++i)
        {
            const auto& filePath = filePaths[i];

            // 获取文件后缀
            std::string extension;
            if (const size_t pos = filePath.find_last_of('.'); pos != std::string::npos)
            {
                extension = filePath.substr(pos);
            }

            if (extension == ".json")
            {
                // 处理JSON文件
                std::cout << "暂不支持JSON文件: " << filePath << std::endl;
            }
            else if (extension == ".geojson" || extension == ".tiff" || extension == ".tif")
            {
                // 记录当前cells的大小
                const size_t prevSize = cells.size();

                // 处理GeoJSON或TIFF文件
                tiffHandler(filePath);

                // 如果是第一个成功加载的文件，保存其元素类型
                if (!typeSet && cells.size() > prevSize)
                {
                    typeSet = true;
                }
            }
            else
            {
                // 不支持的文件类型
                std::cout << "不支持的文件类型: " << extension << "，文件: " << filePath << std::endl;
            }
        }


        std::cout << "========================" << std::endl;
        std::cout << "多文件加载性能报告" << std::endl;
        std::cout << "文件数量: " << filePaths.size() << std::endl;
        std::cout << "加载单元数: " << cells.size() << std::endl;
        std::cout << "========================" << std::endl;
    }
}
_NNU_DGGS_END